#define _CRT_SECURE_NO_WARNINGS

#include <jvmti.h>
#include <jni.h>
#include <string.h>
#include <cstdint>
#include <cstdarg>
#include <iostream>
#include <map>
#include <unordered_map>
#include <vector>
#include <time.h>
#include <Windows.h>
#include <cassert>
#include "Lib.h"
#include "Cache.h"
#include "MethodTable.h"
//#include "java_crw.h"

//#ifdef _DEBUG
#define dprintf printf
//#else
//#define dprintf
//#endif

using namespace std;

//#define checkNull(error,msg) if(error==NULL){printf("[TMon] FATAL ERROR: %s!",msg);exit(-1);}
//#define debug(msg) printf("[TMon]: %s!\n",msg);
//#define debugDetail(msg,code) printf("[TMon]: %s: %d!\n",msg,code);

const char* TrackerPackage = "orca";
const char* TrackerClassName = "orca/OrcaProfilerTracker";
const char* TrackerClassSig = "Lorca/OrcaProfilerTracker;";

const char* TrackerMonitorEnterName = "monitorEnter";
const char* TrackerMonitorEnterSig = "(Ljava/lang/Object;)V";
const char* TrackerMonitorEnterNativeName = "monitorEnter";
const char* TrackerMonitorEnterNativeSig = "(Ljava/lang/Object;Ljava/lang/Object;)V";

const char* LogFileNameClasses = "_classes.orca";
const char* LogFileNameMethods = "_methods.orca";
const char* LogFileNameEvents = "_events.orca";

jvmtiEnv *GJvmti = NULL;
jrawMonitorID GMonitorLock;
jvmtiError error;

bool isLivePhase = false;

struct ThreadInfo
{
	int RelativeDepth;
};

map<jint, ThreadInfo> Threads;

Logger* LogClasses = NULL;
Logger* LogMethods = NULL;
Logger* LogEvents = NULL;

// bool cnt = 0;

bool startWith(const char* s, const char* p)
{
	while (*p)
	{
		if (*s != *p)
			return false;
		s++;
		p++;
	}
	return true;
}

//create synchronized region for each callback
static void enter_critical_section()
{
	error = GJvmti->RawMonitorEnter(GMonitorLock);
	CheckError(error, "RawMonitorEnter");
}

//exit synchronized region
static void exit_critical_section()
{
	error = GJvmti->RawMonitorExit(GMonitorLock);
	CheckError(error, "RawMonitorExit");
}

void DumpCallStack(jthread thread, bool fromOne)
{
	jvmtiFrameInfo finfo[1024];
	jint frameCnt;
	error = GJvmti->GetStackTrace(thread, 0, 1024, finfo, &frameCnt);
	CheckError(error, "GetStackTrace");

	if (error == JVMTI_ERROR_NONE)
	{
		int i = 0;
		if (fromOne)
		{
			i = 1;
			LogEvents->Record(",%d", frameCnt - 1);
		}
		else
			LogEvents->Record(",%d", frameCnt);

		for (; i < frameCnt; i++)
		{
			LogEvents->Record(";%08X;%08X",
				finfo[i].method,
				finfo[i].location);
			// update cache
			CheckCachedMethodSourceInfo(finfo[i].method);
		}
	}
	else
		LogEvents->Record(",0");
}

static void JNICALL callbackMethodEntry(jvmtiEnv *jvmti_env, JNIEnv* env, jthread thread, jmethodID method)
{
	enter_critical_section();
	{
		MethodInfo* minfo = GetCachedMethodInfo(method);
		if (minfo != NULL)
		{
			// append thread stack
			jint threadHash;
			GJvmti->GetObjectHashCode(thread, &threadHash);
			CheckError(error, "GetThreadHashcode_ment");

			// already inside an event, supressed
			if (Threads[threadHash].RelativeDepth >= 0)
				Threads[threadHash].RelativeDepth++;
			// enew event
			else //if (Threads[threadHash].RelativeDepth < 0)
			{
				Threads[threadHash].RelativeDepth = 0;

				// get callee information
				jobject instance = NULL;
				jint instanceHash = 0;
				error = GJvmti->GetLocalInstance(thread, 0, &instance);
				if (error == JVMTI_ERROR_NONE)
					error = GJvmti->GetObjectHashCode(instance, &instanceHash);

				// call method callback
				if (minfo->CallbackEntry)
					minfo->CallbackEntry(thread, threadHash, instance, instanceHash);

			}
		}
	}
	exit_critical_section();
}


//this function will be invoked when a method returns
//it is used to get another part of the implicit context when "notify" exit
static void JNICALL callbackMethodExit(jvmtiEnv *jvmti_env, JNIEnv* env, jthread thread,
	jmethodID method, jboolean was_popped_by_exception, jvalue return_value)
{
	enter_critical_section();
	{
		MethodInfo* minfo = GetCachedMethodInfo(method);
		if (minfo != NULL)
		{
			// append thread stack
			jint threadHash;
			GJvmti->GetObjectHashCode(thread, &threadHash);
			CheckError(error, "GetThreadHashcode");

			// supress nested event
			if (Threads[threadHash].RelativeDepth == 0)
			{
				Threads[threadHash].RelativeDepth = INT_MIN / 2;

				// get callee information
				jobject instance = NULL;
				jint instanceHash = 0;
				error = GJvmti->GetLocalInstance(thread, 0, &instance);
				if (error == JVMTI_ERROR_NONE)
					error = GJvmti->GetObjectHashCode(instance, &instanceHash);

				// call method callback
				if (minfo->CallbackExit)
					minfo->CallbackExit(thread, threadHash, instance, instanceHash);
			}

			// after, dec current frame count
			if (Threads[threadHash].RelativeDepth > 0)
				Threads[threadHash].RelativeDepth--;
		}
	}
	exit_critical_section();
}

static void JNICALL callbackMonitorWait(jvmtiEnv *jvmti_env, JNIEnv* jni_env, jthread thread, jobject object, jlong timeout)
{
	return;
	enter_critical_section();
	{
		jint Ptr = NULL;
		GJvmti->GetObjectHashCode(object, &Ptr);
		printf("[Agent] MonitorWait %.3f, obj hashcode = %d\n", GetTimeSec(), Ptr);
		//gb_jvmti->Deallocate((unsigned char*)&Ptr);
	}
	exit_critical_section();
}

static void JNICALL callbackMonitorWaited(jvmtiEnv *jvmti_env, JNIEnv* jni_env, jthread thread, jobject object, jboolean timed_out)
{
	return;
	enter_critical_section();
	{
		jint Ptr = NULL;
		GJvmti->GetObjectHashCode(object, &Ptr);
		printf("[Agent] MonitorWaited %.3f, obj hashcode = %d\n", GetTimeSec(), Ptr);
		// gb_jvmti->Deallocate((unsigned char*)&Ptr);
	}
	exit_critical_section();
}

static void JNICALL callbackMonitorContendedEnter(jvmtiEnv *jvmti_env, JNIEnv* jni_env, jthread thread, jobject object)
{
	enter_critical_section();
	{
		jint objHash = 0;
		error = GJvmti->GetObjectHashCode(object, &objHash);
		CheckError(error, "GetObjectHashCode");

		jint threadHash = 0;
		error = GJvmti->GetObjectHashCode(thread, &threadHash);
		CheckError(error, "GetObjectHashCode");

		jint ownerHash = 0;
		jvmtiMonitorUsage moInfo;
		error = GJvmti->GetObjectMonitorUsage(object, &moInfo);
		CheckError(error, "GetMonitorUsage");
		if (error == JVMTI_ERROR_NONE)
		{
			if (moInfo.owner != NULL)
			{
				error = GJvmti->GetObjectHashCode(moInfo.owner, &ownerHash);
				CheckError(error, "GetObjectHashCode");
			}
			else
			{
				ownerHash = 0;
			}

			LogEvents->BeginRecord("MonitorContendedEnter,%08X,%08X,%08X", threadHash, objHash, ownerHash);
			DumpCallStack(thread, false);
			LogEvents->EndRecord();

			GJvmti->Deallocate((unsigned char*)moInfo.waiters);
			GJvmti->Deallocate((unsigned char*)moInfo.notify_waiters);
		}
	}
	exit_critical_section();
}

static void JNICALL callbackMonitorContendedEntered(jvmtiEnv *jvmti_env, JNIEnv* jni_env, jthread thread, jobject object)
{
	enter_critical_section();
	{
		jint objHash = 0;
		error = GJvmti->GetObjectHashCode(object, &objHash);
		CheckError(error, "GetObjectHashCode");

		jint threadHash = 0;
		error = GJvmti->GetObjectHashCode(thread, &threadHash);
		CheckError(error, "GetObjectHashCode");

		LogEvents->BeginRecord("MonitorContendedEntered,%08X,%08X", threadHash, objHash);
		DumpCallStack(thread, false);
		LogEvents->EndRecord();
		
	}
	exit_critical_section();
}

static void JNICALL callbackCompiledMethodLoad(jvmtiEnv *jvmti_env, jmethodID method, jint code_size, const void* code_addr, jint map_length, const jvmtiAddrLocationMap* map, const void* compile_info)
{
	return;
	enter_critical_section();
	char* name = NULL;
	char* signature = NULL;
	char* generic_ptr = NULL;
	GJvmti->GetMethodName(method, &name, &signature, &generic_ptr);
//	LogWrite("MethodLoad,%s,%s", name, signature);
	//printf("\nCompiled method load event\n");
	//printf("Method name %s %s %s\n\n", name, signature, generic_ptr);
	
	exit_critical_section();
}

void JNICALL callbackClassLoad(jvmtiEnv *jvmti_env, JNIEnv* jni_env, jthread thread, jclass klass)
{
	enter_critical_section();
	{
		// debug locks
		char* sig = NULL;
		GJvmti->GetClassSignature(klass, &sig, NULL);
		if (strstr(sig, "lock") || strstr(sig, "Lock"))
			dprintf("---lock class %s\n", sig);

		jint methodCount;
		jmethodID* methods;
		error = GJvmti->GetClassMethods(klass, &methodCount, &methods);
		if (error == JVMTI_ERROR_NONE)
		{
			for (int i = 0; i < methodCount; i++)
				TryAddCacheMethodInfo(methods[i], InterestMethodTable);

			GJvmti->Deallocate((unsigned char*)methods);
		}
	}
	exit_critical_section();
}

void JNICALL callbackClassPrepare(jvmtiEnv *jvmti_env, JNIEnv* jni_env, jthread thread, jclass klass)
{
	enter_critical_section();
	{
		jint methodCount;
		jmethodID* methods;
		error = GJvmti->GetClassMethods(klass, &methodCount, &methods);
		if (error == JVMTI_ERROR_NONE)
		{
			for (int i = 0; i < methodCount; i++)
				TryAddCacheMethodInfo(methods[i], InterestMethodTable);

			GJvmti->Deallocate((unsigned char*)methods);
		}
	}
	exit_critical_section();
}

void JNICALL callbackThreadStart(jvmtiEnv *jvmti_env, JNIEnv* jni_env, jthread thread)
{
	enter_critical_section();
	{
		jint threadHash = 0;
		error = GJvmti->GetObjectHashCode(thread, &threadHash);
		CheckError(error, "GetObjectHashCode");

		jvmtiThreadInfo threadInfo;
		error = GJvmti->GetThreadInfo(thread, &threadInfo);

		// initialize thread
		Threads[threadHash].RelativeDepth = INT_MIN / 2;

		const char* threadName = (error == JVMTI_ERROR_NONE) ?
			threadInfo.name : "";

		// dump
		LogEvents->BeginRecord("ThreadStarted,%08X,%s", threadHash, threadName);
		LogEvents->EndRecord();

		if (error == JVMTI_ERROR_NONE)
			GJvmti->Deallocate((unsigned char*)threadInfo.name);	
	}
	exit_critical_section();
}

void JNICALL calllbackThreadEnd(jvmtiEnv *jvmti_env, JNIEnv* jni_env, jthread thread)
{
	enter_critical_section();
	{
		jint threadHash = 0;
		error = GJvmti->GetObjectHashCode(thread, &threadHash);
		CheckError(error, "GetObjectHashCode");

		jlong threadTime = -1;
		error = GJvmti->GetThreadCpuTime(thread, &threadTime);

		// Note here threadTime is actually an unsigned integer.
		uint64_t threadCPUTime = (uint64_t)threadTime;

		// clear threads stack record
		Threads.erase(threadHash);

		LogEvents->BeginRecord("ThreadEnded,%08X,%llu", threadHash, threadCPUTime);
		LogEvents->EndRecord();
	}
	exit_critical_section();
}

static int bci_counter = 0;

void my_crw_fatal_error_handler(const char*message, const char*file, int line)
{
	fprintf(stderr, "mycrwfatalerrorhandler.\n");
}

//void JNICALL callbackClassFileLoadHook(jvmtiEnv *jvmti_env, JNIEnv* jni_env,
//	jclass class_being_redefined, jobject loader,
//	const char* name, jobject protection_domain, jint class_data_len,
//	const unsigned char* class_data, jint* new_class_data_len,
//	unsigned char** new_class_data)
//{
//	enter_critical_section();
//	{
//		const char *classname;
//
//		bci_counter++;
//
//		*new_class_data_len = 0;
//		*new_class_data = NULL;
//
//		/* Name could be NULL */
//		if (name == NULL)
//		{
////			classname = java_crw_demo_classname(class_data, class_data_len, &my_crw_fatal_error_handler);
//			if (classname == NULL)
//				fprintf(stderr, "No classname in classfile");
//		}
//		else
//		{
//			classname = _strdup(name);
//			if (classname == NULL)
//				fprintf(stderr, "Ran out of malloc() space");
//		}
//
//		/* The tracker class itself? */
//		if (strcmp(classname, TrackerClassName) != 0 &&
//			startWith(classname, "java/") == false &&
//			startWith(classname, "sun/") == false)
//		{
//			//ClassIndex            cnum;
//			int                   system_class;
//			unsigned char *       new_image;
//			long                  new_length;
//			int                   len;
//
//			fprintf(stderr, "cbClassFileLoadHook injecting class %s\n", classname);
//
//			/* Is it a system class? */
//			system_class = bci_counter < 8;
//
//			new_image = NULL;
//			new_length = 0;
//
//			/* Call the class file reader/write demo code */
//			java_crw_demo(
//				0,//cnum,
//				classname,
//				class_data,
//				class_data_len,
//				system_class,
//				(char*)TrackerClassName,
//				(char*)TrackerClassSig,
//				NULL,
//				NULL,
//				NULL,
//				NULL,
//				NULL,
//				NULL,
//				(char*)TrackerMonitorEnterName,
//				(char*)TrackerMonitorEnterSig,
//				&new_image,
//				&new_length,
//				&my_crw_fatal_error_handler,
//				NULL);
//
//			if (new_length > 0)
//			{
//				unsigned char *jvmti_space;
//
//				fprintf(stderr, "cbClassFileLoadHook DID inject this class %s.\n", classname);
//				GJvmti->Allocate((jint)new_length, &jvmti_space);
//				(void)memcpy((void*)jvmti_space, (void*)new_image, (int)new_length);
//				*new_class_data_len = (jint)new_length;
//				*new_class_data = jvmti_space; /* VM will deallocate */
//			}
//			else
//			{
//				fprintf(stderr, "cbClassFileLoadHook DID NOT inject this class %s.\n", classname);
//				*new_class_data_len = 0;
//				*new_class_data = NULL;
//			}
//			if (new_image != NULL)
//				free((void*)new_image); /* Free malloc() space with free() */
//		}
//		free((void*)classname);
//
////		LogEvents->BeginRecord("ThreadEnded,%08X,%llu", threadHash, threadCPUTime);
////		LogEvents->EndRecord();
//	}
//	exit_critical_section();
//}

void JNICALL hookedMonitorEnter(JNIEnv *env, jclass c, jobject o)
{
	fprintf(stderr, "Hooked monitor enter!\n");
}

void JNICALL callbackNativeMethodBind(jvmtiEnv *jvmti_env, JNIEnv* jni_env,
	jthread thread, jmethodID method, void* address, void** new_address_ptr)
{
	enter_critical_section();
	{
		char* name;
		char* sig;

		int error = GJvmti->GetMethodName(method, &name, &sig, NULL);
		CheckError(error, "GetObjectHashCode");

		fprintf(stderr, "NaiveBInding.... %s\n", name);

		if (strcmp("monitorEnter", name) == 0)
		{
			*new_address_ptr = &hookedMonitorEnter;
		}
		GJvmti->Deallocate((unsigned char*)name);
		GJvmti->Deallocate((unsigned char*)sig);
	}
	exit_critical_section();
}

void JNICALL callbackVMStart(jvmtiEnv *jvmti_env, JNIEnv* jni_env)
{
	dprintf("VMSTart----------------------------------\n");
}

void JNICALL callbackVMDeath(jvmtiEnv *jvmti_env, JNIEnv* jni_env)
{
	dprintf("VMDEATH----------------------------------\n");
}

void JNICALL callbackVMInit(jvmtiEnv *jvmti_env, JNIEnv* jni_env, jthread thread)
{
//	dprintf("VMInit----------------------------------\n");
	error = GJvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_CLASS_LOAD, (jthread)NULL);
	error = GJvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_CLASS_PREPARE, (jthread)NULL);

	error = GJvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_METHOD_ENTRY, (jthread)NULL);
	error = GJvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_METHOD_EXIT, (jthread)NULL);
	error = GJvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_MONITOR_CONTENDED_ENTER, (jthread)NULL);
	error = GJvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_MONITOR_CONTENDED_ENTERED, (jthread)NULL);
	error = GJvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_THREAD_START, (jthread)NULL);
	error = GJvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_THREAD_END, (jthread)NULL);
//	error = GJvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_CLASS_FILE_LOAD_HOOK, (jthread)NULL);
//	error = GJvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_NATIVE_METHOD_BIND, (jthread)NULL);
}

//when JVM init, this function is invoked
JNIEXPORT jint JNICALL Agent_OnLoad(JavaVM *jvm, char *options, void *reserved)
{
	jvmtiEventCallbacks callbacks;

	jint result = jvm->GetEnv((void**)&GJvmti, JVMTI_VERSION_1_2);

	if (result != JNI_OK || GJvmti == NULL)
	{
		printf("ERROR: Unable to access JVMTI!\n");
		return JNI_ERR;
	}

	//enable the capabilities of jvmti
	jvmtiCapabilities gb_capa;
	memset(&gb_capa, 0, sizeof(jvmtiCapabilities));

	//gb_capa.can_signal_thread = 1;
	//gb_capa.can_get_owned_monitor_info = 1;
	//gb_capa.can_generate_method_exit_events = 1;
	//gb_capa.can_generate_method_entry_events = 1;
	//gb_capa.can_generate_exception_events = 1;
	//gb_capa.can_generate_vm_object_alloc_events = 1;
	//gb_capa.can_tag_objects = 1;
	//gb_capa.can_generate_all_class_hook_events = 1;
	//gb_capa.can_generate_native_method_bind_events = 1;
	//gb_capa.can_access_local_variables = 1;
	//gb_capa.can_get_monitor_info = 1;
	//gb_capa.can_generate_monitor_events = 1;
	//gb_capa.can_get_owned_monitor_info = 1;
	GJvmti->GetPotentialCapabilities(&gb_capa);

	error = GJvmti->AddCapabilities(&gb_capa);
	if (error != JVMTI_ERROR_NONE)
	{
		printf("ERROR: Can't set JVMTI capabilities\n");
		return JNI_ERR;
	}

	// register the callbacks
	memset(&callbacks, 0, sizeof(jvmtiEventCallbacks));

	callbacks.MethodEntry = &callbackMethodEntry;
	callbacks.MethodExit = &callbackMethodExit;
//	callbacks.MonitorWait = &callbackMonitorWait;
//	callbacks.MonitorWaited = &callbackMonitorWaited;
	callbacks.MonitorContendedEnter = &callbackMonitorContendedEnter;
	callbacks.MonitorContendedEntered = &callbackMonitorContendedEntered;
	callbacks.ThreadStart = &callbackThreadStart;
	callbacks.ThreadEnd = &calllbackThreadEnd;
	callbacks.ClassLoad = &callbackClassLoad;
	callbacks.ClassPrepare = &callbackClassPrepare;
//	callbacks.ClassFileLoadHook = &callbackClassFileLoadHook;
//	callbacks.NativeMethodBind = &callbackNativeMethodBind;

//	callbacks.VMStart = &callbackVMStart;
	callbacks.VMInit = &callbackVMInit;
//	callbacks.VMDeath = &callbackVMDeath;

	error = GJvmti->SetEventCallbacks(&callbacks, sizeof(callbacks));
	if (error != JVMTI_ERROR_NONE)
	{	printf("ERROR: Can't set jvmti callback!\n");
		return JNI_ERR;
	}

	error = GJvmti->CreateRawMonitor("Aura", &GMonitorLock);
	if (error != JVMTI_ERROR_NONE)
	{
		printf("ERROR: Can't create raw monitor!\n");
		return JNI_ERR;
	}

	//register the event notification for each event
	error = GJvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_VM_START, (jthread)NULL);
	error = GJvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_VM_INIT, (jthread)NULL);
	error = GJvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_VM_DEATH, (jthread)NULL);

	error = GJvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_CLASS_LOAD, (jthread)NULL);
	error = GJvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_CLASS_PREPARE, (jthread)NULL);

	// open logfile
	char buf[1024];
	strcpy(buf, "I");
	LogClasses = Logger::CreateLogger(strcat(buf, LogFileNameClasses));
	strcpy(buf, "I");
	LogMethods = Logger::CreateLogger(strcat(buf, LogFileNameMethods));
	strcpy(buf, "I");
	LogEvents = Logger::CreateLogger(strcat(buf, LogFileNameEvents));

	// init featured method table
	InitInterestMethodTable();

	// OK
	LogEvents->BeginRecord("AgentLoaded");
	LogEvents->EndRecord();

	return JNI_OK;

}

JNIEXPORT void JNICALL Agent_OnUnload(JavaVM *vm)
{
	LogEvents->BeginRecord("AgentUnLoad");
	LogEvents->EndRecord();

	delete LogClasses;
	delete LogMethods;
	delete LogEvents;
}
